import pprint, copy, math, tarfile, os, sys, gzip, cPickle, time
import numpy as np
from pywolb.deterministic import base
from pywolb.utilities import diversetools, numpytools, stringtools
#from pywolb.utilities import animationtools
#import matplotlib.pyplot as plt

for module in [base,diversetools,numpytools,stringtools]:
    reload(module)

expand = numpytools.expand

def build_M(self, migration_rates):
        migration_rates = diversetools.substitute(migration_rates, np.nan, 0.)
        M = np.array( migration_rates )
        r = int( math.sqrt(len(M)) )
        M.shape = (r,r)
        for k in range(r):
            M[k,k] = 1 - np.nansum(M[k])
        return M

class MetaPopulation(object):    
    def __init__(self, migration_rates, *populations):
        """Usage: migration_rates= numpy array
                  populations = pop1, pop2, ...
        Attention: nan throws a pickle error in python <2.5, use 0. instead.
        """
        self.populations = populations
        self.number_of_populations = len(populations)
        for i in range(self.number_of_populations):
            self.__dict__["pop%d" % (i+1)] = self.populations[i]
        self.generation = self.pop1.generation
        self._pop_ndim = self.pop1._ndim
        self._pop_shape = self.pop1._data_shape
        self._build_migration(migration_rates)
        self.eq_precision = 1e-7
        self.eq = self._eq_reached(self.eq_precision)
        self.target_reached = None
        self.hybridzoneshape = ()
    
    def show(self, *args, **kwargs):
        print self._show(*args, **kwargs)
    
    def _show(self, *args, **kwargs):
        s = ''
        if args:
            if args[0] == 'overview':
                return self._show_overview()
            elif args[0] == 'full overview':
                try:
                    n_most_frequent = args[1]
                except:
                    n_most_frequent = 5
                return self._show_overview(fullmode=True,n_most_frequent=n_most_frequent)
            elif args[0] == 'parameters':
                s = stringtools.two_columns('migration:', str(self.M))
                s += '\n'
        i = 1
        for pop in self.populations:
            s += "pop. %d:\n" % i
            s += pop._show(*args, **kwargs)
            s += '\n\n'
            i += 1
        return s.strip()
    
#    def show_overview(self):
#        l = []
#        for i in range(self.number_of_populations):
#            pop = self.__dict__["pop%d" % (i+1)] 
#            s = pop._overview2string()
#            longest = stringtools.max_line(s)
#            s = "pop. %d".center(longest) % (i+1) + '\n' + \
#                '-'*longest + '\n' + \
#                s
#            l.append(s)
#        print "generation  %d\n" % self.generation
#        print stringtools.columns(l, num_spaces=3)
    
    def show_overview(self, fullmode=False, n_most_frequent=5):
        print self._show_overview(fullmode, n_most_frequent)
    
    def _show_overview(self, fullmode=False, n_most_frequent=5):
        returnstring = "generation %d\n" % self.generation
        l = []
        for i in range(self.number_of_populations):
            pop = self.__dict__["pop%d" % (i+1)] 
            if fullmode:
                s = pop._full_overview2string(n_most_frequent)
            else:
                s = pop._overview2string()
            longest = stringtools.max_line(s)
            s = "pop. %d".center(longest) % (i+1) + '\n' + \
                '-'*longest + '\n' + \
                s
            l.append(s)
        returnstring += stringtools.columns(l, num_spaces=3)
        return returnstring.strip()
    
    def visualize(self, loci='all', generation=-1, withlabels=True, showaxes=False, **kwargs):
        animationtools.metapopviz(self, loci=loci, generation=generation, withlabels=withlabels, **kwargs)
    
    def _get_max_change(self):
        change = np.array( [pop._get_max_change() for pop in self.populations] )
        return np.max(change)
    
    def reset(self):
        self.generation = 0
        for pop in self.populations:
            pop.reset()
        for i in range(self.number_of_populations):
            self.__dict__["pop%d" % (i+1)] = self.populations[i]
    
    def save(self, filename):
        """Save everything in human readable text format."""
        outfile = open(filename, 'w')
        print >> outfile, "%s Overview %s\n" % ('#'*10,'#'*20), self._show('overview'), '\n\n', \
                          "%s Parameters %s\n" % ('#'*10,'#'*20), self._show('parameters'), '\n\n', \
                          "%s Frequencies %s\n" % ('#'*10,'#'*20), self._show(), '\n\n', \
                          "%s Weights %s\n" % ('#'*10,'#'*20), self._show('weights')
        outfile.close()
    
    def _state2string(self, *types):
        max_change = self._get_max_change()
        freqs = [self.get_sum(t) for t in types]
        s = "%13.3g" % max_change
        for f in freqs:
            s += "%10.3g" % f
        return s
               
    def _outputline(self, *types):
        print '\r',
        s = "%10d" % self.generation
        for pop in self.populations:
            s += '  ' + pop._state2string(*types)
        sys.stdout.write(s)
        sys.stdout.flush()
        sys.stdout.flush()
        
    def _shortoutputline(self, *types):
        print '\r',
        s = "%10d" % self.generation
        for pop in self.populations:
            s += '  ' + pop._shortstate2string(*types)
        sys.stdout.write(s)
        sys.stdout.flush()
        sys.stdout.flush()
    
    def _output_title1(self, *types):
        s = "%10s" % ''
        i = 1
        for pop in self.populations:
            name = "pop%d" % i
            name_width = len(name)
            total_width = len(pop._state2string(*types))
            d,m = divmod(total_width-name_width-2, 2)
            i += 1
            s += "  %s %s %s" % (d*'-', name, (d+m)*'-')
        return s
    
    def _output_title2(self, *types):
        s = "%10s" % 'generation'
        for pop in self.populations:
            s += '  ' + pop._state_title(*types)
        return s
    
    def _output_short_title1(self, *types):
        s = "%10s" % ''
        i = 1
        for pop in self.populations:
            name = "pop%d" % i
            name_width = len(name)
            total_width = len(pop._shortstate2string(*types))
            d,m = divmod(total_width-name_width-2, 2)
            i += 1
            s += "  %s %s %s" % (d*'-', name, (d+m)*'-')
        return s
    
    def _output_short_title2(self, *types):
        s = "%10s" % 'generation'
        for pop in self.populations:
            s += '  ' + pop._state_short_title(*types)
        return s
    
    def _build_migration(self, migration_rates):
        self.migration_rates = migration_rates
        self.M = self.migration_rates
        self.M_expanded = expand(migration_rates, 2+self._pop_ndim, [0,1])
    
    def add_population(self, pop):
        if pop._data_shape == self._pop_shape:
            self.populations.append(pop)
            self.number_of_populations += 1
            self.__dict__["p%d" % ( self.number_of_populations + 1 )] = pop
        else:
            raise ValueError, "shape mismatch: population must have shape %s" % str(self._pop_shape)
    
    def set_migration(self, migration_rates):
        self._build_migration(migration_rates)
    
    def _copy_data(self):
        # converting a population into an array automatically creates a copy:
        return np.array( [np.array(pop) for pop in self.populations] )
    
    def _migration(self):
        # get population frequencies:
        frequencies = self._copy_data()
        r = self._pop_ndim
        frequencies_expanded = expand( frequencies, r+2, range(1,r+2) )
        # the actual migration step:
        frequencies = np.sum( frequencies_expanded * self.M_expanded, 1 )
        # set population frequencies:
        for i in range(self.number_of_populations):
            self.populations[i].set_array( frequencies[i] )
            self.populations[i].normalize()
        
    def _next_generation(self):
        for pop in self.populations:
            pop._bookkeeping_before_cycle()
        self._migration()
        for pop in self.populations:
            pop._next_generation_within_metapop()
        self.generation += 1
        for pop in self.populations:
            pop._bookkeeping_after_cycle()
    
    def _hstore(self):
        for pop in self.populations:
            pop._hstore()
    
    def _threshold_hstore(self, threshold=1e-2):
        """Store current state if any frequency in any population has (totally) changed by more than threshold
since the last storage."""
        store_required = False
        # check whether in any populations, frequencies have changed by more than the threshold value:
        for pop in self.populations:
#            if pop._get_total_change_since_last_hstore() >= threshold:
            if np.any(pop._get_changes_since_last_hstore() >= threshold):
                store_required = True
                break
        # if so, than store each population:
        if store_required:
            self._hstore()
    
    def run(self, number_of_generations, output_modulo=20, output_types=None, output_mode='long', threshold=1e-2, timeit=True):
        if timeit:
            t0 = time.time()
        for pop in self.populations:
            pop._build_constant_reproduction_factor()
            pop._build_constant_viability_selection_factor()
            #pop._set_repro_func()
            pop._bind_methods()
        try:
            if output_types:
                if output_mode == 'long':
                    print self._output_title1(*output_types)
                    print self._output_title2(*output_types)
                    for i in range(number_of_generations):
                        self._next_generation()
                        if self.generation % output_modulo == 0:
                            self._outputline(*output_types)
                        self._threshold_hstore(threshold)
                        if numpytools.ipshell_available:
                            numpytools.ipbreak()
                elif output_mode == 'short':
                    print self._output_short_title1(*output_types)
                    print self._output_short_title2(*output_types)
                    for i in range(number_of_generations):
                        self._next_generation()
                        if self.generation % output_modulo == 0:
                            self._shortoutputline(*output_types)
                        self._threshold_hstore(threshold)
                        if numpytools.ipshell_available:
                            numpytools.ipbreak()
            else:
                for i in range(number_of_generations):
                    self._next_generation()
                    self._threshold_hstore(threshold)
                    if numpytools.ipshell_available:
                        numpytools.ipbreak()
        finally:
            self._hstore()
            for pop in self.populations:
                pop._detach_methods()
            if timeit:
                print "\nruntime: %.1f seconds" % (time.time()-t0)
    
    def run_hstore(self, number_of_generations, store_interval=1, output_modulo=20, output_types=None, timeit=True):
        """Run and store every store_interval generations.
        
        Beware of huge data accumulation! This is primarily meant to be used for debugging and for measuring effective migration rates."""
        if timeit:
            t0 = time.time()
        for pop in self.populations:
            pop._build_constant_reproduction_factor()
            pop._build_constant_viability_selection_factor()
            #pop._set_repro_func()
            pop._bind_methods()
        try:
            if output_types:
                print self._output_title1(*output_types)
                print self._output_title2(*output_types)
                for i in range(number_of_generations):
                    self._next_generation()
                    if self.generation % output_modulo == 0:
                        self._outputline(*output_types)
                    if self.generation % store_interval == 0:
                        self._hstore()
                    if numpytools.ipshell_available:
                        numpytools.ipbreak()
            else:
                for i in range(number_of_generations):
                    self._next_generation()
                    if self.generation % store_interval == 0:
                        self._hstore()
                    if numpytools.ipshell_available:
                        numpytools.ipbreak()
        finally:
            self._hstore()
            for pop in self.populations:
                pop._detach_methods()
            if timeit:
                print "\nruntime: %.1f seconds" % (time.time()-t0)
    
    def runto(self, nstop, output_modulo=20, output_types=None, threshold=1e-2, timeit=True):
        """Run until generation nstop is reached."""
        n = nstop - self.generation
        self.run(n, output_modulo=output_modulo, output_types=output_types, threshold=threshold, timeit=timeit)
    
    def runtofreq(self, target_type='P1', target_freq='0.5', target_pop=0, nmin=10, nmax=5000, precision=1e-7, additive=True, threshold=1e-2):
        """Run simulation until target_type has reached a given frequency target_freq in 
        population target_pop but stop when an equilibrium has been reached or after nmax generations.
        This method can be used to measure rates of spread.
        """
        if additive:
            nmax += self.generation
        if nmin > 0:
            self.run(nmin, threshold=threshold)
        for pop in self.populations:
            pop._build_constant_reproduction_factor()
            pop._build_constant_viability_selection_factor()
            #pop._set_repro_func()
            pop._bind_methods()
        try:
            while (self.populations[target_pop].get_sum(target_type) < target_freq) and (self.generation < nmax) and (not self._eq_reached(precision)):
                self._next_generation()
                self._threshold_hstore(threshold)
                if numpytools.ipshell_available:
                    numpytools.ipbreak()
        finally:
            self._hstore()
            for pop in self.populations:
                pop._detach_methods()
            self.eq = self._eq_reached(precision)
            self.target_reached = (self.populations[target_pop].get_sum(target_type) >= target_freq)
            return self.target_reached

    def _eq_reached(self, precision=1e-7):
        for pop in self.populations:
            if not pop._eq_reached(precision):
                return False
        return True
                
    def findEQ(self, precision=1e-7, nmin=10, nmax=5e5, additive=True, interval=20, output_types=None, output_mode='long', threshold=1e-2, timeit=True):
        """Find EQ frequencies.
        Run until all individual frequencies in all populations do not change more than 
        precision between consecutive generations. Run for at least nmin generations, 
        and stop simulation after nmax generation even if no EQ has been reached.
        """
        if timeit:
            t0 = time.time()
        self.eq_precision = precision
        self.eq = self._eq_reached(precision)
        if self.eq:
            return True
        if additive:
            nmax += self.generation
        self.run(nmin, threshold=threshold,timeit=False) #, output_modulo=output_modulo, output_types=output_types)
        for pop in self.populations:
            pop._build_constant_reproduction_factor()
            pop._build_constant_viability_selection_factor()
            #pop._set_repro_func()
            pop._bind_methods()
        try:
            if output_types:
                if output_mode == 'long':
                    print self._output_title1(*output_types)
                    print self._output_title2(*output_types)
                    while (not self._eq_reached(precision)) and (self.generation < nmax):
                        for i in range(interval):
                            self._next_generation()
                            self._threshold_hstore(threshold)
                            if numpytools.ipshell_available:
                                numpytools.ipbreak()
                        self._outputline(*output_types)
                elif output_mode == 'short':
                    print self._output_short_title1(*output_types)
                    print self._output_short_title2(*output_types)
                    for i in range(number_of_generations):
                        self._next_generation()
                        if self.generation % output_modulo == 0:
                            self._shortoutputline(*output_types)
                        self._threshold_hstore(threshold)
                        if numpytools.ipshell_available:
                            numpytools.ipbreak()
                else:
                    print 'something wrong with output_mode'
            else:
                while (not self._eq_reached(precision)) and (self.generation < nmax):
                    for i in range(interval):
                        self._next_generation()
                        self._threshold_hstore(threshold)
                        if numpytools.ipshell_available:
                            numpytools.ipbreak()
#        except:
#            raise                
        finally:
            if output_types:
                print
            self._hstore()
            for pop in self.populations:
                pop._detach_methods()
            self.eq = self._eq_reached(precision)
            return self.eq
            if timeit:
                print "\nruntime: %.1f seconds" % (time.time()-t0)
    
    def findIncrease(self, pop='pop1', allele='P1', nmin=100, repetition=5, output_types=None, debug=False):
        self.run(nmin)
        freqs = [None for i in range(repetition)]
        increases = [None for i in range(repetition)]
        for p in self.populations:
            p._build_constant_reproduction_factor()
            p._build_constant_viability_selection_factor()
            #p._set_repro_func()
            p._bind_methods()
        try:
            if output_types:
                # long output mode:
                print self._output_title1(*output_types)
                print self._output_title2(*output_types)
                for i in range(repetition):
                    self._next_generation()
                    self._hstore()
                    freqs[i] = self.__dict__[pop].get_sum(allele)
                    increases[i] = self.__dict__[pop]._allele_increases(allele)
                    if numpytools.ipshell_available:
                        numpytools.ipbreak()
                self._outputline(*output_types)
            else:
                for i in range(repetition):
                    self._next_generation()
                    self._hstore()
                    freqs[i] = self.__dict__[pop].get_sum(allele)
                    increases[i] = self.__dict__[pop]._allele_increases(allele)
                    if numpytools.ipshell_available:
                        numpytools.ipbreak()
#        finally:
        finally:
            if output_types:
                print
            self._hstore()
            for p in self.populations:
                p._detach_methods()
            if debug:
                print "%s frequencies (%s): %s" % (allele, pop, str(freqs))
                print "increases:", increases
            return increases[-1]  # np.alltrue( increases )
        
    def findEQ_intervalstore(self, precision=1e-7, nmin=10, nmax=5e5, additive=True, store_interval=1, interval=20, output_types=None, output_mode='long'):
        """Find EQ frequencies.
        Run until all individual frequencies in all populations do not change more than 
        precision between consecutive generations. Run for at least nmin generations, 
        and stop simulation after nmax generation even if no EQ has been reached.
        """
        self.eq_precision = precision
        self.eq = self._eq_reached(precision)
        if self.eq:
            return True
        if additive:
            nmax += self.generation
        self.run_hstore(nmin, store_interval=store_interval) #, output_modulo=output_modulo, output_types=output_types)
        for pop in self.populations:
            pop._build_constant_reproduction_factor()
            pop._build_constant_viability_selection_factor()
            #pop._set_repro_func()
            pop._bind_methods()
        try:
            if output_types:
                if output_mode == 'long':
                    print self._output_title1(*output_types)
                    print self._output_title2(*output_types)
                    while (not self._eq_reached(precision)) and (self.generation < nmax):
                        for i in range(interval):
                            self._next_generation()
                            if self.generation % store_interval == 0:
                                self._hstore()
                            if numpytools.ipshell_available:
                                numpytools.ipbreak()
                        self._outputline(*output_types)
                elif output_mode == 'short':
                    print self._output_short_title1(*output_types)
                    print self._output_short_title2(*output_types)
                    for i in range(number_of_generations):
                        self._next_generation()
                        if self.generation % output_modulo == 0:
                            self._shortoutputline(*output_types)
                        if self.generation % store_interval == 0:
                            self._hstore()
                        if numpytools.ipshell_available:
                            numpytools.ipbreak()
                else:
                    print 'something wrong with output_mode'
            else:
                while (not self._eq_reached(precision)) and (self.generation < nmax):
                    for i in range(interval):
                        self._next_generation()
                        if self.generation % store_interval == 0:
                            self._hstore()
                        if numpytools.ipshell_available:
                            numpytools.ipbreak()
#        except:
#            raise                
        finally:
            if output_types:
                print
            self._hstore()
            for pop in self.populations:
                pop._detach_methods()
            self.eq = self._eq_reached(precision)
            return self.eq

    
    def _get_P1_tendency(self):
        """Use only with 2 populations!"""
        # P1 frequencies in the two populations
        p1 = []
        p1.append( self.pop1.get_sum('P1') )
        p1.append( self.pop2.get_sum('P1') )
        p1 = np.array(p1)
#        diff = np.abs( p1[0]-p1[1] )
        # same thing for the last stored state
        p1prev = []
        p1f = base.LabeledArray(self.pop1._previous_frequencies, self.pop1.labels)
        p2f = base.LabeledArray(self.pop2._previous_frequencies, self.pop2.labels)
        p1prev.append( p1f.get_sum('P1') )
        p1prev.append( p2f.get_sum('P1') )
        p1prev = np.array( p1prev )
#        if diff >= precision:
#            return 'unresolved'
#        else:
        if np.alltrue( p1-p1prev > 0 ):
            return 'increasing'
        elif np.alltrue( p1-p1prev < 0 ):
            return 'decreasing'
        else:
            return 'unresolved'
    
    def _get_P1_tendency_h(self):
        """Use only with 2 populations!"""
        h1 = self.pop1.history.get_sum('P1')
        h2 = self.pop2.history.get_sum('P1')
        # P1 frequencies in the two populations
        p1 = []
        p1.append( h1[-1] )
        p1.append( h2[-1] )
        p1 = np.array(p1)
#        diff = np.abs( p1[0]-p1[1] )
        # same thing for the last stored state
        p1prev = []
        p1prev.append( h1[-2] )
        p1prev.append( h2[-2] )
        p1prev = np.array( p1prev )
#        if diff >= precision:
#            return 'unresolved'
#        else:
#        print 'diff:', np.array2string(p1-p1prev, precision=10)
        if np.alltrue( p1-p1prev > 0 ):
            return 'increasing'
        elif np.alltrue( p1-p1prev < 0 ):
            return 'decreasing'
        else:
            return 'unresolved'

    def findTendency(self, nmin=10, nmax=1e5, additive=True, interval=20, output_types=None, threshold=1e-2):
        """Find tendency of P1 development: decrease or increase, i.e., extinction or spread.
        precision applies to the difference in P1 frequency between the populations
        """
        self._tendency = 'unresolved'
        if additive:
            nmax += self.generation
        self.run(nmin, threshold=threshold) #, output_modulo=output_modulo, output_types=output_types)            
#        self.run(nmin/2, threshold=threshold) #, output_modulo=output_modulo, output_types=output_types)
#        self.run(nmin/2, threshold=threshold)     # make sure that hstore is called twice
        for pop in self.populations:
            pop._build_constant_reproduction_factor()
            pop._build_constant_viability_selection_factor()
            #pop._set_repro_func()
            pop._bind_methods()
        try:
            if output_types:
                print self._output_title1(*output_types)
                print self._output_title2(*output_types)
                while (self._get_P1_tendency_h() == 'unresolved') and (self.generation < nmax):
                    for i in range(interval):
                        self._next_generation()
                        self._threshold_hstore(threshold)
                        if numpytools.ipshell_available:
                            numpytools.ipbreak()
                    self._outputline(*output_types)
            else:
                while (self._get_P1_tendency_h() == 'unresolved') and (self.generation < nmax):
                    for i in range(interval):
                        self._next_generation()
                        self._threshold_hstore(threshold)
                        if numpytools.ipshell_available:
                            numpytools.ipbreak()
#        except:
#            raise                
        finally:
            if output_types:
                print
            self._hstore()
            for pop in self.populations:
                pop._detach_methods()
            self._tendency = self._get_P1_tendency()
            return self._tendency

    def findEQstoreFreqs(self, target_type, precision=1e-7, nmin=10, nmax=5e5, additive=True, threshold=1e-2):
        """Find EQ frequencies.
        Run until all individual frequencies in all populations do not change more than 
        precision between consecutive generations. Run for at least nmin generations, 
        and stop simulation after nmax generation even if no EQ has been reached.
        """
        freqs_store = []
        for i in range(self.number_of_populations):
            freqs_store.append([])
        self.eq_precision = precision
        self.eq = self._eq_reached(precision)
        if self.eq:
            return True, freqs_store
        if additive:
            nmax += self.generation
#        self.run(nmin, threshold=threshold) #, output_modulo=output_modulo, output_types=output_types)
        for pop in self.populations:
            pop._build_constant_reproduction_factor()
            pop._build_constant_viability_selection_factor()
            #pop._set_repro_func()
            pop._bind_methods()
        try:
            while (self.generation < nmin) or ((not self._eq_reached(precision)) and (self.generation < nmax)):
                self._next_generation()
                self._threshold_hstore(threshold)
                for i in range(self.number_of_populations):
                    x = self.populations[i].get_sum(target_type)
                    freqs_store[i].append( x )
                if numpytools.ipshell_available:
                    numpytools.ipbreak()
        finally:
            self._hstore()
            for pop in self.populations:
                    pop._detach_methods()
            self.eq = self._eq_reached(precision)
            return self.eq, np.array(freqs_store)
    
    def advance(self, generations):
        self.generation += generations
        for pop in self.populations:
            pop.generation_increase(generations)
    
    def introduce_allele(self, pop, allele_name, frequency):
        self.__dict__[pop].introduce_allele(allele_name, frequency, advance=False)
        self.advance(1)
    
    def introduce_alleles(self, *args):
        """args are tuples (population, allele, frequency)"""
        n = len(args)/3
        for i in range(n):
            pop, allele_name, frequency = args[i*3:(i+1)*3]
            self.__dict__[pop].introduce_allele(allele_name, frequency, advance=False)
        self.advance(1)
        
    def introduce_type(self, pop, _type, freq):
        self.__dict__[pop]._set_type_frequency(_type, freq, advance=False)
        self.advance(1)
    
    def _create_focal_population(self, pop, allele):
        """Usage within metapopulation: focal_pop = _create_focal_population(self.pop1, 'np1')"""
        import weights4
        neutral_alleles = ['np1','np2','np3']
        labels = copy.deepcopy(pop.labels)
        labels['axes'].append('neutral')
        labels['elements'].append(neutral_alleles)
        sh = self._pop_shape + (3,)
        freqs = np.zeros(sh)
        allele_index = neutral_alleles.index(allele)
        freqs[...,allele_index] = copy.deepcopy(pop.data)
        focal_pop = base.FocalPopulation(freqs, labels)
        focal_pop.add_weights( *pop.get_weights() )
        focal_pop.add_weight(weights4.NuclearInheritanceAtSingleLocus('neutral', 'np', range(1,4)))
        return focal_pop
    
    def _create_background_population(self, pop):
        """Usage within metapopulation: background_pop = _create_background_population(self.pop1)"""
        import weights4
        neutral_alleles = ['np1','np2','np3']
        labels = copy.deepcopy(pop.labels)
        labels['axes'].append('neutral')
        labels['elements'].append(neutral_alleles)
        sh = self._pop_shape + (3,)
        freqs = np.zeros(sh)
        freqs[...,2] = copy.deepcopy(pop.data)      # background populations all have 'np3'!
        background_pop = base.BackgroundPopulation(freqs, labels)
        background_pop.add_weights( *pop.get_weights() )
        background_pop.add_weight(weights4.NuclearInheritanceAtSingleLocus('neutral', 'np', range(1,4)))
        return background_pop
    
    def _create_effective_metapopulation(self, i, j):
        npops = self.number_of_populations
        pops = ['dummy_pop' for k in range(npops)]
        pops[i] = self._create_focal_population(self.populations[i], 'np1')
        pops[j] = self._create_focal_population(self.populations[j], 'np2')
        for k in range(npops):
            if k not in [i,j]:
                pops[k] = self._create_background_population(self.populations[k])
        metapop = MetaPopulation(self.M, *pops)
        return metapop
    
    def _calculate_effective_migration_matrix(self, n=100, allele='np1', debug=False):
        """Run all pairwise population combinations as focal populations with the other populations as background populations to gain the effective migration matrix of the metapopulation."""
        if debug:
            global metapop
        effM = np.nan * np.ones(np.shape(self.M))
        npops = self.number_of_populations
        for i in range(npops):
            for j in range(i+1,npops):
                metapop = self._create_effective_metapopulation(i,j)
                if debug:
                    p = i+1
                    q = j+1
                    print '#'*80
                    print "effective migration for population pair pop%d - pop%d:" % (p,q)
                    metapop.show('overview')
                    metapop.run_hstore(1)
                    metapop.show('overview')
                metapop.run_hstore(n)
                effM[i,j],effM[j,i] = effective_migration(metapop.populations[i], metapop.populations[j], 'np1')
                if debug:
                    metapop.show('overview')                    
                    print "    m%d%d = %f" % (p,q,effM[i,j])
                    print "    m%d%d = %f" % (q,p,effM[j,i])
        return effM
    
    def calculate_effective_migration_matrix(self, n=100, allele='np1', debug=False):
        self.effective_migration = self._calculate_effective_migration_matrix(n,allele,debug)
        self.effM = self.effective_migration
    
    def calculate_gene_flow_factors(self):
        if not hasattr(self,'effM'):
            self.calculate_effective_migration_matrix()
        self.GFF = self.effM / self.M
    
    def get_effective_migration_rates2(self, n=100, allele='np1'):
        """Use only for two population metapopulations!
        n - number of generations to run with neutral locus added"""
        if self.number_of_populations not in [2]:
            raise ValueError, 'method only applicable to metapopulations consisting of two populations'
        import weights4
        labels = copy.deepcopy(self.pop1.labels)
        labels['axes'].append('neutral')
        labels['elements'].append(['np1','np2'])
        sh = self._pop_shape + (2,)
        freqs1 = np.zeros(sh)
        freqs2 = np.zeros(sh)
        freqs1[...,0] = self.pop1.data
        freqs2[...,1] = self.pop2.data
        population1 = base.Population(freqs1, labels)
        for t in self.pop1.weights:
            for w in self.pop1.weights[t].values():
                population1.add_weight(copy.deepcopy(w))
        population1.add_weight(weights4.NuclearInheritanceAtSingleLocus('neutral', 'np', range(1,3)))
        population2 = base.Population(freqs2, labels)
        for t in self.pop2.weights:
            for w in self.pop2.weights[t].values():
                population2.add_weight(copy.deepcopy(w))
        population2.add_weight(weights4.NuclearInheritanceAtSingleLocus('neutral', 'np', range(1,3)))
        metapop = MetaPopulation(self.migration_rates, population1, population2)
        metapop.run_hstore(n)
        return effective_migration2(metapop.pop1, metapop.pop2, allele)
    
    def get_effective_migration_rates(self, n=100, allele='np1'):
        """Use only for two population metapopulations!
        n - number of generations to run with neutral locus added"""
        numpops = self.number_of_populations
        if numpops not in [2,3]:
            raise ValueError, 'method only applicable to metapopulations consisting of two or three populations'
        import weights4
        labels = copy.deepcopy(self.pop1.labels)
        labels['axes'].append('neutral')
        labels['elements'].append(["np%d" % (i+1) for i in range(numpops)])
        sh = self._pop_shape + (numpops,)
        freqs = []
        for i in range(numpops):
            freqs.append(np.zeros(sh))
            code = "freqs[%d][...,%d] = self.pop%d.data" % (i,i,i+1)
            exec code
        pops = []
        for i in range(numpops):
            pops.append(base.Population(freqs[i], labels))
            for t in self.populations[i].weights:
                for w in self.populations[i].weights[t].values():
                    pops[i].add_weight(copy.deepcopy(w))
            pops[i].add_weight(weights4.NuclearInheritanceAtSingleLocus('neutral', 'np', range(1,numpops+1)))
        metapop = MetaPopulation(self.migration_rates, *pops)
        print "metapopulation, iterating to get effective migration rates..."
        metapop.run_hstore(n)
        metapop.show_overview()
        if numpops == 2:
            return effective_migration2(metapop.pop1, metapop.pop2, allele)
        else:
            return effective_migration3(metapop.pop1, metapop.pop2, metapop.pop3, allele)
    
    def get_meff(self):
        return self.get_effective_migration_rates()
    
    def divergence_at(self, locus, precision=1e-5):
        """For a 2 population metapopulation, return True if the two populations diverge at the locus, and return False if they do not."""
        return not np.alltrue( 
                np.abs(
                    np.array(self.pop1.get_sum(locus)) - np.array(self.pop2.get_sum(locus))
                    ) < precision )
    
    def get_sum(self, _type):
        result = []
        for i in range(self.number_of_populations):
            result.append(self.populations[i].get_sum(_type))
        return np.array(result)
    
    def get_hsum(self, _type):
        result = []
        for i in range(self.number_of_populations):
            result.append(self.populations[i].history.get_sum(_type))
        result = np.array(result)
        return result.swapaxes(0,1)   # axes of actual returned result are 'generation', 'population'
    
    def get_generations(self):
        return np.array(self.pop1.history.generations)


def effective_migration(pop1, pop2, allele):
    """See 'The concept of effective migration rate and its application in speciation models'
    by Yutaka Kobayashi and Arndt Telschow"""
    pop1prev = base.Population(pop1._previous_frequencies, pop1.labels)
    pop2prev = base.Population(pop2._previous_frequencies, pop2.labels)
    p = [pop1prev.get_sum(allele), pop1.get_sum(allele)]        # [p(t), p(t+1)]
    q = [pop2prev.get_sum(allele), pop2.get_sum(allele)]        # [q(t), q(t+1)]
    t = 0
    m1 = (p[t+1]-p[t]) / (q[t]-p[t])
    m2 = (q[t+1]-q[t]) / (p[t]-q[t])
    return np.nan_to_num(m1), np.nan_to_num(m2)

#def effective_migration3(pop1, pop2, pop3, allele):
#    """This function unfortunaltely does not make sense!"""
#    labels = pop1.labels
#    pop1t1 = base.Population(pop1.history[-2], labels)       # previous freqs
#    pop1t2 = base.Population(pop1.history[-3], labels)       # previous to previous freqs
#    pop2t1 = base.Population(pop2.history[-2], labels)       # previous freqs
#    pop2t2 = base.Population(pop2.history[-3], labels)       # previous to previous freqs
#    pop3t1 = base.Population(pop3.history[-2], labels)       # previous freqs
#    pop3t2 = base.Population(pop3.history[-3], labels)       # previous to previous freqs
#    x1 =[pop1t2.get_sum(allele), pop1t1.get_sum(allele), pop1.get_sum(allele)]
#    x2 =[pop2t2.get_sum(allele), pop2t1.get_sum(allele), pop2.get_sum(allele)]
#    x3 =[pop3t2.get_sum(allele), pop3t1.get_sum(allele), pop3.get_sum(allele)]
#    t = 0
#    # the following code was calculated with Mathematica:
#    m12 = (math.pow(x1[t+1],2) + x1[t+2]*x3[t] + x1[t]*(-x1[t+2] + x3[t+1]) - x1[t+1]*(x3[t] + x3[t+1])) / (x1[t+1]*(x2[t] - x3[t]) + x2[t+1]*x3[t] - x2[t]*x3[t+1] + x1[t]*(-x2[t+1] + x3[t+1]))
#    m13 = (math.pow(x1[t+1],2) + x1[t+2]*x2[t] + x1[t]*(-x1[t+2] + x2[t+1]) - x1[t+1]*(x2[t] + x2[t+1])) / (-x2[t+1]*x3[t] + x1[t+1]*(-x2[t] + x3[t]) + x1[t]*(x2[t+1] - x3[t+1]) + x2[t]*x3[t+1])
#    m21 = (x1[t+1]*x2[t+1] - x1[t]*x2[t+2] - x1[t+1]*x3[t] + x2[t+2]*x3[t] + x1[t]*x3[t+1] - x2[t+1]*x3[t+1]) / (x1[t+1]*x2[t] - x1[t]*x2[t+1] - x1[t+1]*x3[t] + x2[t+1]*x3[t] + x1[t]*x3[t+1] - x2[t]*x3[t+1])
#    m23 = (x1[t+1]*x2[t] - x1[t]*x2[t+1] - x1[t+1]*x2[t+1] + math.pow(x2[t+1],2) + x1[t]*x2[t+2] - x2[t]*x2[t+2]) / (x1[t+1]*x2[t] - x1[t]*x2[t+1] - x1[t+1]*x3[t] + x2[t+1]*x3[t] + x1[t]*x3[t+1] - x2[t]*x3[t+1])
#    m31 = (x1[t+1]*(x3[t] - x3[t+1]) + math.pow(x3[t+1],2) - x3[t]*x3[t+2] + x1[t]*(-x3[t+1] + x3[t+2])) / (-x2[t+1]*x3[t] + x1[t+1]*(-x2[t] + x3[t]) + x1[t]*(x2[t+1] - x3[t+1]) + x2[t]*x3[t+1])
#    m32 = (x1[t+1]*x2[t] - x1[t]*x2[t+1] - x1[t+1]*x3[t+1] + x2[t+1]*x3[t+1] + x1[t]*x3[t+2] - x2[t]*x3[t+2]) / (x1[t+1]*x2[t] - x1[t]*x2[t+1] - x1[t+1]*x3[t] + x2[t+1]*x3[t] + x1[t]*x3[t+1] - x2[t]*x3[t+1])
#    return m12,m13, m21,m23, m31,m32
    
def main():
    pass
    

    
    
